import numpy as np
import matplotlib.pyplot as plt

Nx = 64
Ny = 64
L = 1 # 求解域长度
dx = 1 / Nx
dy = 1 / Ny



# 泊松方程等式右边项
f = np.zeros((Nx,Ny))
rsq = np.zeros((Nx,Ny))
for i in range(Nx):
    for j in range(Ny):
        rsq[i,j] = (i*dx - L/2)**2 + (j*dx - L/2)**2
        sigsq = 0.08
        f[i,j] = np.exp(-rsq[i,j] / (2*sigsq))*(rsq[i,j] - 2*sigsq)/(sigsq**2)



f_hat = np.fft.fft2(f)

k = np.zeros((Nx))
l = np.zeros((Nx))
for i in range(0,Nx):
    if i < Nx/2:
        l[i] = k[i] = 2*np.pi/L*i
    else:
        l[i] = k[i] = 2*np.pi/L*(i - Nx)

# 泊松等式左边项
lhs = np.zeros((Nx,Ny),dtype = complex)
for i in range(Nx):
    for j in range(Ny):
        lhs[i,j] = - (k[i]**2 + l[j]**2)
lhs[0,0] = 1# 防止被0除

u_hat = f_hat / lhs
u = np.fft.ifft2(u_hat).real # 得到正确的结果，也就是压力泊松方程的压力
u -= u[0,0]

# 检验结果，用解析及和数值解比较
analytic = np.zeros((Nx,Ny))
error = np.zeros((Nx,Ny)) 
for i in range(1,Nx-1):
    for j in range(1,Ny-1):
        analytic[i,j] = np.exp(-rsq[i,j] / (2 * sigsq))
        error[i,j] = abs(analytic[i,j] - u[i,j])

max_error=np.max(error)
print(max_error)
# 为了绘图，我们需要将解析解和数值解限制在相同的范围内
# 并且移除可能由于FFT导致的高频噪声造成的误差
u_analytic = analytic - np.mean(analytic)  # 移除解析解的均值
u_numeric = u.real - np.mean(u.real)      # 移除数值解的均值

# 绘制数值解和解析解
plt.figure(figsize=(12, 6))

# 第一个子图：数值解
plt.subplot(1, 2, 1)
plt.title("Numerical Solution")
plt.imshow(u_numeric, extent=[0, L, 0, L], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')

# 第二个子图：解析解
plt.subplot(1, 2, 2)
plt.title("Analytical Solution")
plt.imshow(u_analytic, extent=[0, L, 0, L], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')

# 调整子图间距
plt.tight_layout()

# 显示误差
plt.figure()
plt.title("Error between Numerical and Analytical Solutions")
plt.imshow(error, extent=[0, L, 0, L], origin='lower', cmap='coolwarm')
plt.colorbar(label='Error')

# 显示所有图形
plt.show()